from template import *



import matplotlib.pyplot as plt
import numpy as np

# Fixing random state for reproducibility
np.random.seed(19680801)

nParticle = 3000

vIncident = np.zeros(3)
vIncident[0] = 1.0e6
vIncident[1] = 1.0e6
vIncident[2] = 1.0e6

momentum_unit = np.zeros(3)
momentum_unit[0] = vIncident[0] / np.linalg.norm(vIncident)
momentum_unit[1] = vIncident[1] / np.linalg.norm(vIncident)
momentum_unit[2] = vIncident[2] / np.linalg.norm(vIncident)

ke = 1000.0

pi = 3.1415926


vx = np.zeros(nParticle)
vy = np.zeros(nParticle)
vz = np.zeros(nParticle)

for i in range(0, nParticle):
    ra = np.random.rand()
    cosX = (2.0 + ke - 2.0 * math.pow(1.0+ke,ra)) / ke
    sinX = math.sqrt(1.0 - cosX*cosX)

    ra = np.random.rand()
    phi = 2.0 * pi *ra
    cosPhi = math.cos(phi)
    sinPhi = math.sin(phi)

    ve = np.linalg.norm(vIncident)

    r11 = momentum_unit[0]
    r12 = momentum_unit[1]
    r13 = momentum_unit[2]

    up1 = 0.0
    up2 = 0.0
    up3 = 0.0
    if r11 == 1.0 or r11 == -1.0:
        up1 = 0.0
        up2 = 1.0
        up3 = 0.0
    else:
        up1 = 1.0
        up2 = 0.0
        up3 = 0.0
    
    r21 = r12 * up3 - r13 * up2
    r22 = r13 * up1 - r11 * up3
    r23 = r11 * up2 - r12 * up1

    mag = math.sqrt(r21*r21 + r22*r22 + r23*r23)

    r21 = r21 / mag
    r22 = r22 / mag
    r23 = r23 / mag

    r31 = r22 * r13 - r23 * r12
    r32 = r23 * r11 - r21 * r13
    r33 = r21 * r12 - r22 * r11



    vx[i] = ve * (r11 * cosX + r21 * sinX * sinPhi + r31 * sinX * cosPhi)
    vy[i] = ve * (r12 * cosX + r22 * sinX * sinPhi + r32 * sinX * cosPhi)
    vz[i] = ve * (r13 * cosX + r23 * sinX * sinPhi + r33 * sinX * cosPhi)
    print(vx[i], vy[i], vz[i])


fig = plt.figure()
ax = fig.add_subplot(projection='3d')


ax.scatter(vx, vy, vz)

ax.set_xlabel('vx Label')
ax.set_ylabel('vy Label')
ax.set_zlabel('vz Label')

plt.show()